
#include <windows.h>
#include <wtsapi32.h>
#include <tchar.h>
#include <stdio.h>
#include <iostream>
#include <vcclr.h>


#include "Stdafx.h"
#include "SessionLauncher.h"

using namespace System;

DWORD Spawn(wchar_t * appName, wchar_t *cmd, bool spawnInConsoleSession, bool waitExit)
{
	CHandle job(CreateJobObject(nullptr, nullptr));
	if (!job) {
		return  GetLastError();
		return 0;
	}

	CHandle ioPort(CreateIoCompletionPort(INVALID_HANDLE_VALUE, nullptr, 0, 1));


	if (!ioPort) {
		return GetLastError();
	}

	JOBOBJECT_ASSOCIATE_COMPLETION_PORT port;
	port.CompletionKey = job;
	port.CompletionPort = ioPort;
	if (!SetInformationJobObject(job, JobObjectAssociateCompletionPortInformation, &port, sizeof(port))) {
		return  GetLastError();
	}

	
	JOBOBJECT_BASIC_LIMIT_INFORMATION basicLimit;
	basicLimit.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
	JOBOBJECT_EXTENDED_LIMIT_INFORMATION extendedLimit;
	extendedLimit.BasicLimitInformation = basicLimit;

	if (!SetInformationJobObject(job, JobObjectExtendedLimitInformation, &extendedLimit, sizeof(extendedLimit))) {
		return GetLastError();
	}


	PROCESS_INFORMATION processInformation;
	STARTUPINFO startupInfo = { sizeof(startupInfo) };
	startupInfo.dwFlags = STARTF_USESHOWWINDOW;
	startupInfo.wShowWindow = SW_HIDE;

//	wchar_t *dsktp = L"WinSta0\\Default"; //allow process to be interactive
	//startupInfo.lpDesktop = dsktp;

	if (!spawnInConsoleSession) {
		if (!CreateProcess(nullptr, cmd, nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr, nullptr, &startupInfo, &processInformation)) {
			return GetLastError();
		}
	}
	else {
		HANDLE svctoken = nullptr;
		HANDLE provtoken = nullptr;
		DWORD sessionId = WTSGetActiveConsoleSessionId();

		// Create token for new process by duplicating our own token
		if (OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &svctoken) == 0)
		{
			return GetLastError();
		}
		DuplicateTokenEx(svctoken, TOKEN_ALL_ACCESS, nullptr, SecurityImpersonation, TokenPrimary, &provtoken);
		CloseHandle(svctoken);
		svctoken = nullptr;
		// Move the new token to the target session
		if (!SetTokenInformation(provtoken, TokenSessionId, &sessionId, sizeof(DWORD)))
		{
			return GetLastError();
		}

		// Set up the new process information
		if (!CreateProcessAsUser(provtoken, appName, cmd, nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr, nullptr, &startupInfo, &processInformation))
		{
			return GetLastError();
		}

		if (provtoken)
		{
			CloseHandle(provtoken);
		}
	}
	//if this fails the spawned process remain suspended !!!!
	if (!AssignProcessToJobObject(job, processInformation.hProcess)) {
		return GetLastError();
	}

	ResumeThread(processInformation.hThread);
	CloseHandle(processInformation.hThread);
	CloseHandle(processInformation.hProcess);


	if (waitExit) {
		DWORD completionCode;
		ULONG_PTR completionKey;
		LPOVERLAPPED overlapped;
		while (GetQueuedCompletionStatus(ioPort, &completionCode, &completionKey, &overlapped, INFINITE)
			&& !((HANDLE)completionKey == job
			&& completionCode == JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO)) {
			//wait
		}
	}

	//all ok
	return 0;
}
